#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
微调效果评估脚本
用于测试微调前后模型在特定领域知识上的表现
"""

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import argparse

# 测试问题集合
TEST_QUESTIONS = [
    "什么是量子计算？",
    "请解释量子纠缠现象",
    "什么是区块链技术？",
    "解释一下机器学习中的过拟合现象",
    "什么是Transformer架构？",
    "介绍一下BERT模型的特点",
    "什么是联邦学习？",
    "解释一下边缘计算的概念",
    "什么是元宇宙？",
    "请说明5G网络的关键技术"
]

def load_model(model_path):
    """
    加载模型和分词器
    
    Args:
        model_path (str): 模型路径
    """
    print(f"正在加载模型: {model_path}")
    
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
    
    # 加载模型
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="balanced",
        offload_folder="./offload",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        local_files_only=True
    )
    
    return model, tokenizer

def generate_response(model, tokenizer, prompt, max_length=200):
    """
    生成模型回复
    
    Args:
        model: 模型
        tokenizer: 分词器
        prompt (str): 输入提示
        max_length (int): 最大生成长度
    """
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    inputs = inputs.to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_length=max_length,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # 移除输入提示部分，只返回生成的部分
    response = generated_text[len(prompt):].strip()
    
    return response

def evaluate_model(model_path):
    """
    评估模型表现
    
    Args:
        model_path (str): 模型路径
    """
    # 加载模型
    model, tokenizer = load_model(model_path)
    
    print(f"\n正在评估模型: {model_path}")
    print("=" * 50)
    
    # 对每个问题生成回答
    for i, question in enumerate(TEST_QUESTIONS, 1):
        print(f"\n问题 {i}: {question}")
        try:
            response = generate_response(model, tokenizer, question)
            print(f"回答: {response}")
        except Exception as e:
            print(f"生成回答时出错: {e}")
        print("-" * 30)

def main():
    parser = argparse.ArgumentParser(description="评估微调模型效果")
    parser.add_argument(
        "--model_path",
        default="./models/Qwen1.5-1.8B",  # 原始模型路径
        help="模型路径"
    )
    parser.add_argument(
        "--finetuned_model_path",
        help="微调后模型路径"
    )
    
    args = parser.parse_args()
    
    # 评估原始模型
    print("评估原始模型:")
    evaluate_model(args.model_path)
    
    # 如果提供了微调后模型路径，则也评估微调后模型
    if args.finetuned_model_path and os.path.exists(args.finetuned_model_path):
        print("\n" + "=" * 60)
        print("评估微调后模型:")
        evaluate_model(args.finetuned_model_path)
        
        print("\n" + "=" * 60)
        print("对比分析:")
        print("通过对比原始模型和微调后模型的回答，可以评估微调效果。")
        print("微调后的模型应该在专业领域知识方面有更好的表现。")

if __name__ == "__main__":
    main()